import torch
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader


def load_cifar10_data(batch_size, img_size, root='../data'):
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(img_size),
        # transforms.RandomCrop(32, padding=4),  # 先四周填充0，在吧图像随机裁剪成32*32
        transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转 选择一个概率概率
        transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])  # 标准化
    ])

    # 加载数据集
    train_dataset = datasets.CIFAR10(root=root, train=True, download=True, transform=transform)
    train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [int(len(train_dataset)*0.8), int(len(train_dataset)*0.2)])
    test_dataset = datasets.CIFAR10(root=root, train=False, download=True, transform=transform)

    # 创建DataLoader
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    return train_loader, val_loader, test_loader, train_dataset.dataset.classes